import pandas as pd
import os
import sys
import time
import random
import concurrent.futures
from threading import Lock
from tqdm import tqdm
import threading
from concurrent.futures import ThreadPoolExecutor

# 添加项目根目录到系统路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入谷歌翻译工具和代理管理器
from utils.google_translate import translate_from_ch_to_en, translate_from_en_to_cn
from proxy_manager import init_proxy_pool

# 创建一个锁用于同步进度条更新
progress_lock = Lock()

def process_single_text(text, use_proxy=True, max_retries=3, sleep_time_range=(0.5, 2.0), timeout=20):
    """
    处理单个文本的数据增强
    
    参数:
        text: 要增强的中文文本
        use_proxy: 是否使用代理
        max_retries: 最大重试次数
        sleep_time_range: 随机延迟的时间范围(最小值, 最大值)
        timeout: 请求超时时间(秒)
    
    返回:
        增强后的文本，如果处理失败则返回None
    """
    if not text or not isinstance(text, str) or text.strip() == "":
        return None
    
    # 随机延迟，避免请求过于频繁
    time.sleep(random.uniform(sleep_time_range[0], sleep_time_range[1]))
    
    try:
        # 中文 -> 英文
        english_translation = translate_from_ch_to_en(
            text, 
            use_proxy=use_proxy, 
            max_retries=max_retries,
            timeout=timeout
        )
        if not english_translation:
            return None
        
        # 再次随机延迟
        time.sleep(random.uniform(sleep_time_range[0], sleep_time_range[1]))
        
        # 英文 -> 中文
        augmented_text = translate_from_en_to_cn(
            english_translation, 
            use_proxy=use_proxy, 
            max_retries=max_retries,
            timeout=timeout
        )
        return augmented_text
    except Exception as e:
        print(f"处理文本时出错: {str(e)}")
        return None

def _worker_process_batch(texts, use_proxy=True, max_retries=3, sleep_time_range=(0.5, 2.0), timeout=20):
    """
    工作线程处理一批文本
    
    参数:
        texts: 要处理的文本列表
        use_proxy: 是否使用代理
        max_retries: 最大重试次数
        sleep_time_range: 随机延迟的时间范围
        timeout: 请求超时时间(秒)
    
    返回:
        处理结果的列表，与输入文本列表一一对应
    """
    results = []
    for text in texts:
        augmented_text = process_single_text(
            text, 
            use_proxy=use_proxy, 
            max_retries=max_retries,
            sleep_time_range=sleep_time_range,
            timeout=timeout
        )
        results.append(augmented_text)
    return results

def augment_data_with_translation(
    input_data, 
    text_column="text", 
    output_file=None,
    sample_ratio=1.0, 
    random_seed=42, 
    sleep_time_range=(0.5, 2.0), 
    num_workers=4,
    batch_size=5,
    use_proxy=True,
    max_retries=3,
    timeout=20
):
    """
    使用翻译方法增强数据
    
    参数:
        input_data: 输入数据，可以是CSV文件路径或DataFrame对象
        text_column: 文本列的名称
        output_file: 输出CSV文件路径，如果为None则不保存文件
        sample_ratio: 采样比例，范围[0,1]
        random_seed: 随机种子
        sleep_time_range: 请求之间随机延迟的时间范围(最小值, 最大值)
        num_workers: 工作线程数量
        batch_size: 每批处理的数据量
        use_proxy: 是否使用代理
        max_retries: 最大重试次数
        timeout: 请求超时时间(秒)
    
    返回:
        增强后的DataFrame
    """
    # 根据输入类型处理数据
    if isinstance(input_data, str):
        print(f"开始数据增强处理，输入文件: {input_data}")
        # 读取CSV文件
        df = pd.read_csv(input_data)
    elif isinstance(input_data, pd.DataFrame):
        print("开始数据增强处理，输入DataFrame")
        df = input_data.copy()
    else:
        raise TypeError("input_data 必须是CSV文件路径或DataFrame对象")
    
    print(f"原始数据大小: {len(df)}行")
    
    # 采样处理
    if sample_ratio < 1.0:
        random.seed(random_seed)
        df = df.sample(frac=sample_ratio, random_state=random_seed)
        print(f"采样后数据大小: {len(df)}行 (采样比例: {sample_ratio})")
    
    # 准备处理的文本
    texts D:\Alvin\code\paper_code = df[text_column].fillna("").tolist()
    augmented_texts = [None] * len(texts)
    
    # 分批处理数据
    batches = []
    for i in range(0, len(texts), batch_size):
        batches.append(texts[i:i+batch_size])
    
    print(f"使用 {num_workers} 个工作线程处理 {len(batches)} 批数据")
    
    # 使用线程池执行批处理任务
    with tqdm(total=len(batches), desc="正在增强数据") as pbar:
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            # 提交所有批次任务
            futures = []
            for batch_idx, batch in enumerate(batches):
                future = executor.submit(
                    _worker_process_batch, 
                    batch, 
                    use_proxy, 
                    max_retries, 
                    sleep_time_range,
                    timeout
                )
                future.batch_idx = batch_idx
                future.start_idx = batch_idx * batch_size
                futures.append(future)
            
            # 收集结果
            for future in concurrent.futures.as_completed(futures):
                batch_results = future.result()
                start_idx = future.start_idx
                
                # 将批次结果合并到总结果中
                for i, result in enumerate(batch_results):
                    if start_idx + i < len(augmented_texts):
                        augmented_texts[start_idx + i] = result
                
                pbar.update(1)
    
    # 创建结果列并输出文件
    result_df = df.copy()
    result_df['augmented_text'] = augmented_texts
    
    # 添加失败统计信息
    success_count = sum(1 for text in augmented_texts if text is not None)
    fail_count = len(augmented_texts) - success_count
    success_rate = success_count / len(augmented_texts) * 100 if augmented_texts else 0
    
    print(f"数据增强完成，成功: {success_count}条，失败: {fail_count}条，成功率: {success_rate:.2f}%")
    
    # 保存结果到文件（如果指定了输出文件）
    if output_file:
        print(f"保存增强结果到: {output_file}")
        result_df.to_csv(output_file, index=False)
    
    return result_df

def combine_original_and_augmented(df, text_column="text", augmented_column="augmented_text"):
    """
    合并原始文本和增强文本
    
    参数:
        df: 包含原始文本和增强文本的DataFrame
        text_column: 原始文本列名
        augmented_column: 增强文本列名
    
    返回:
        合并后的DataFrame
    """
    # 创建新的DataFrame，只保留成功增强的行
    valid_rows = df[df[augmented_column].notnull()].copy()
    
    # 将原始文本和增强文本分别加入不同行
    combined_rows = []
    
    for _, row in valid_rows.iterrows():
        # 添加原始文本行
        original_row = row.copy()
        original_row['is_augmented'] = 0
        combined_rows.append(original_row)
        
        # 添加增强文本行
        augmented_row = row.copy()
        augmented_row[text_column] = row[augmented_column]
        augmented_row['is_augmented'] = 1
        combined_rows.append(augmented_row)
    
    # 创建新的DataFrame
    combined_df = pd.DataFrame(combined_rows)
    
    # 丢弃增强文本列
    if augmented_column in combined_df.columns:
        combined_df = combined_df.drop(columns=[augmented_column])
    
    return combined_df

def main():
    """
    主函数，用于命令行调用
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='通过翻译链路进行数据增强')
    parser.add_argument('--input', type=str, required=True, help='输入CSV文件路径')
    parser.add_argument('--output', type=str, default=None, help='输出CSV文件路径')
    parser.add_argument('--sample_ratio', type=float, default=1.0, help='采样比例，默认为1.0（处理所有数据）')
    parser.add_argument('--random_seed', type=int, default=42, help='随机种子，默认为42')
    parser.add_argument('--min_sleep', type=float, default=0.5, help='最小休眠时间（秒）')
    parser.add_argument('--max_sleep', type=float, default=1.5, help='最大休眠时间（秒）')
    parser.add_argument('--combine', action='store_true', help='是否将增强数据与原始数据合并')
    parser.add_argument('--num_workers', type=int, default=4, help='线程数，默认为4')
    
    args = parser.parse_args()
    
    # 进行数据增强
    augmented_df = augment_data_with_translation(
        input_data=args.input,
        output_file=args.output,
        sample_ratio=args.sample_ratio,
        random_seed=args.random_seed,
        sleep_time_range=(args.min_sleep, args.max_sleep),
        num_workers=args.num_workers
    )
    
    # 如果需要合并原始数据和增强数据
    if args.combine:
        if args.output is None:
            base_name = os.path.splitext(os.path.basename(args.input))[0]
            augmented_file = os.path.join(os.path.dirname(args.input), f"{base_name}_augmented.csv")
        else:
            augmented_file = args.output
        
        combine_original_and_augmented(
            df=augmented_df
        )

if __name__ == "__main__":
    main() 